import streamlit as st
from langchain_openai import ChatOpenAI
import os
# 定义模型


#定义左侧边栏
with st.sidebar:
    st.markdown(f"""
       <center>
           <img src="https://vip.helloimg.com/1/2024/07/02/66841f6f4a3a5.png"  width="100" />
           <h1>Chatbot</h1>
       </center>    
    """, unsafe_allow_html=True)

    # 角色定义输入框
    system_message = st.text_area("角色定义", value="你是一个能够帮助用户的AI助手.")

    # 创造力调节Temperature
    temperature = st.slider("创造力调节", min_value=0.0, max_value=2.0, value=0.7, step=0.1, format="%.1f" , help="值越大，越有创造力")

#定义 右边的chatbot对话窗口
st.title(" AI 聊天机器人")

#初始化界面的聊天列表
if "messages" not in st.session_state:
    st.session_state.messages = [{
        "role": "assistant",
        "content": "你好，我是一个AI助手，你可以问我任何问题。"
    }]

if "message_history" not in st.session_state:
    message_history = []

# 显示聊天消息
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

llm = ChatOpenAI(model_name="deepseek-chat"
                 , api_key=os.environ["DEEPSEEK_API_KEY"]
                 ,base_url=os.environ["DEEPSEEK_BASE_URL"]
                 ,temperature=temperature)
def chat_stream(user_input,system_message):
    if system_message is not None:
        message_history.append({"role": "system", "content": system_message})
    message_history.append({"role": "user", "content": user_input})

    return llm.stream(message_history)



# 获取用户输入
if user_input := st.chat_input("请输入你的问题"):
    # 显示用户输入
    st.chat_message("user").markdown(user_input)

    # 将用户输入添加到聊天列表
    st.session_state.messages.append({"role": "user", "content": user_input})

    with st.chat_message("assistant"):
        with st.spinner(""):
            response = chat_stream(user_input, system_message)
            # 创建显示消息的容器
            message_placeholder = st.empty()

            # Ai的答案
            ai_response = ""
            for chunk in response:
                ai_response += chunk.content
                message_placeholder.markdown(ai_response + " ")

            # 在聊天窗口输出完整的AI答案
            message_placeholder.markdown(ai_response)

            message_history.append({"role": "assistant", "content": ai_response})

            # 将AI的答案添加到聊天列表
            st.session_state.messages.append({"role": "assistant", "content": ai_response})